package top.wanzc.java.basic.concurrent;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.*;

/**
 * @className:TransferMDCThreadPoolDemo 
 * @comment: 注释信息 
 * @date:     Nov 16, 2025 10:24:10 AM 
 * @author   zhengchong.wan
 * @since    V1.0.0	 
 */
class MDC {
	public static void clear() {
		System.out.println("clear MDC : " + Thread.currentThread().getId());
	}
	public static void setContextMap(Map<String, String> contextMap) {
		System.out.println("set context map : " + Thread.currentThread().getId());
	}
	public static Map<String, String> getCopyOfContextMap() {
		System.out.println("get context map : " + Thread.currentThread().getId());
		Map<String, String> map = new HashMap<String, String>();
		map.put("threadId", "1");
		return map;
	}
}

class MDCThreadPoolExecutor extends ThreadPoolExecutor {
	
    public MDCThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit,
    		BlockingQueue<Runnable> workQueue) {
		super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
	}

	@Override
    public void execute(Runnable task) {
        //super.execute(TransferMDCThreadPoolDemo.wrap(task, MDC.getCopyOfContextMap()));
		super.execute(task);
    }

    @Override
    public Future<?> submit(Runnable task) {
        return super.submit(TransferMDCThreadPoolDemo.wrap(task, MDC.getCopyOfContextMap()));
    }

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        return super.submit(TransferMDCThreadPoolDemo.wrap(task, MDC.getCopyOfContextMap()));
    }
}


public class TransferMDCThreadPoolDemo {
	
	public static void main(String[] args) {
		//method1();
		method2();
	}
	
	private static void method1() {
		System.out.println("------------------method1--------------------");
		ExecutorService exec = Executors.newCachedThreadPool();
		Runnable task = () -> System.out.println("do some task : " + Thread.currentThread().getId());
		exec.submit(wrap(task, MDC.getCopyOfContextMap()));
		exec.shutdown();
	}
	
	private static void method2() {
		System.out.println("------------------method2--------------------");
		ThreadPoolExecutor threadPoolExecutor = new MDCThreadPoolExecutor(
		3, 5, 1, TimeUnit.MINUTES, new ArrayBlockingQueue<Runnable>(1000)
		);
        threadPoolExecutor.submit(() -> System.out.println("do some task : " + Thread.currentThread().getId()));
        threadPoolExecutor.shutdown();
	}
	
	public static <T> Callable<T> wrap(final Callable<T> callable,
			final Map<String, String> context) {
		// context代表需要在线程间传递的MDC中存储的key/value
		return () -> {
			if (context == null || context.size() == 0) {
				MDC.clear();
			} else {
				MDC.setContextMap(context);
			}
			try {
				return callable.call();
			} finally {
				MDC.clear();
			}
		};
	}
	
	public static Runnable wrap(final Runnable runnable,
			final Map<String, String> context) {
		return () -> {
			if (context == null || context.size() == 0) {
				MDC.clear();
			} else {
				MDC.setContextMap(context);
			}
			try {
				runnable.run();
			} finally {
				MDC.clear();
			}
		};
	}

}

